import json
import re
import time

from openai import OpenAI


# 获取回答
def get_completion(client, assistant_id, thread_id, user_input, funcs, debug=False):
    """
    Executes a completion request with the given parameters.

    Args:
        assistant_id (str): The ID of the assistant.
        thread_id (str): The ID of the thread.
        user_input (str): The user input content.
        funcs (list): A list of functions.
        debug (bool, optional): Whether to print debug information. Defaults to False.

    Returns:
        str: The message as a response to the completion request.
    """
    if debug:
        print("获取回答...")

    # 创建 Message
    message = client.beta.threads.messages.create(
        thread_id=thread_id,
        role="user",
        content=user_input
    )

    # 创建 Run
    run = client.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id=assistant_id,
    )

    # 运行 Run
    while True:
        while run.status in ['queued', 'in_progress']:
            run = client.beta.threads.runs.retrieve(
                thread_id=thread_id,
                run_id=run.id
            )
            time.sleep(1)

        # 执行 function
        if run.status == "requires_action":
            tool_calls = run.required_action.submit_tool_outputs.tool_calls
            tool_outputs = []
            for tool_call in tool_calls:
                if debug:
                    print(str(tool_call.function))
                func = next(iter([func for func in funcs if func.__name__ == tool_call.function.name]))
                try:
                    output = func(**eval(tool_call.function.arguments))
                except Exception as e:
                    output = "Error: " + str(e)

                if debug:
                    print(f"{tool_call.function.name}: ", output)

                tool_outputs.append(
                    {
                        "tool_call_id": tool_call.id,
                        "output": json.dumps(output)
                    }
                )

            run = client.beta.threads.runs.submit_tool_outputs(
                thread_id=thread_id,
                run_id=run.id,
                tool_outputs=tool_outputs
            )
        elif run.status == "failed":
            raise Exception("Run Failed. Error: ", run.last_error)
        else:
            messages = client.beta.threads.messages.list(
                thread_id=thread_id
            )
            message = messages.data[0].content[0].text.value
            pattern = r"/imgs/\d{10}\.png"
            match = re.search(pattern, message)
            if match:
                message = {"image": match.group()}
            if debug:
                print(message)
            return message


if __name__ == '__main__':
    client = OpenAI(api_key="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", timeout=10000)
    my_assistant = client.beta.assistants.retrieve("asst_4H8LNBta8d2wcKOzETQZ8yfB")
    print(my_assistant)

    # Thread(id='thread_SgLCtadALcweSFwV8BGwAmiO', created_at=1701316784, metadata={}, object='thread')
    # empty_thread = client.beta.threads.create()
    my_thread = client.beta.threads.retrieve("thread_SgLCtadALcweSFwV8BGwAmiO")
    print(my_thread)

    # hread_message = client.beta.threads.messages.create(
    #     my_thread.id,
    #     role="user",
    #     content="Go语言hello world 怎么写"
    # )
    # print(hread_message)

    # run = client.beta.threads.runs.create(
    #     thread_id=my_thread.id,
    #     assistant_id=my_assistant.id
    # )
    # print(run)
    #
    # run_steps = client.beta.threads.runs.steps.list(
    #     thread_id=my_thread.id,
    #     run_id=run.id
    # )
    # print(run_steps)

    get_completion(client,my_assistant.id,my_thread.id,"我早上吃了一碗蛋炒饭，请问这包含多少热量",[],True)


